# 导入所需的库
import cv2  # OpenCV库，用于图像处理和二维码检测
import pandas as pd  # Pandas库，用于数据处理和DataFrame操作
from hashdata import hash_dataframe  # 导入哈希函数，用于数据完整性验证
import numpy as np  # NumPy库，用于数值计算
import json  # JSON库，用于解析二维码中的JSON数据
import time  # 时间模块，用于时间控制和计时
from typing import List, Dict, Optional  # 类型提示
from pathlib import Path  # 路径处理
from io import StringIO  # 字符串IO操作，用于解析CSV数据
from pymongo import MongoClient  # MongoDB客户端库，用于数据库操作

class QRCodeDataReconstructor:
    """
    二维码数据重建器类
    
    该类使用OpenCV的WeChatQRCode库进行二维码扫描解码，并根据解码后的数据分块
    重新构建完整的数据框(DataFrame)。支持从多张图片或摄像头实时扫描二维码。
    支持时间段检查功能，可根据设定的工作时间自动启停。
    """
    
    def __init__(self, model_dir: str = "models"):
        """
        初始化WeChatQRCode检测器
        
        参数:
            model_dir: 模型文件目录路径，默认值为"models"
        """
        # 将模型目录字符串转换为Path对象
        self.model_dir = Path(model_dir)
        # 初始化检测器为None
        self.detector = None
        # 调用初始化检测器方法
        self.initialize_detector()
        
        # 默认工作时间设置（早8点到晚6点）
        self.work_start_time = (8, 0)  # 小时, 分钟
        self.work_end_time = (18, 0)   # 小时, 分钟
        self.is_in_work_time = True    # 默认处于工作时间模式
        
        # MongoDB连接信息
        self.mongo_ip = '10.0.0.233'  # 工控网IP 解码侧
        self.SVR_MONGO_URI = 'mongodb://eCDUsvr:eCDUsvr2020@{0}:27017'.format(self.mongo_ip)
        self.SVR_DB_NAME = 'eCDU'
        
    def initialize_detector(self) -> bool:
        """
        初始化WeChatQRCode检测器
        
        加载WeChatQRCode所需的模型文件并创建检测器实例
        包括检测模型和超分辨率模型各两个文件
        
        返回:
            bool: 初始化是否成功
        """
        try:
            # 初始化WeChatQRCode检测器，需要四个模型文件
            # detect.prototxt和detect.caffemodel用于二维码检测
            # sr.prototxt和sr.caffemodel用于提高二维码图像分辨率
            self.detector = cv2.wechat_qrcode_WeChatQRCode(
                str(self.model_dir / "detect.prototxt"),
                str(self.model_dir / "detect.caffemodel"),
                str(self.model_dir / "sr.prototxt"), 
                str(self.model_dir / "sr.caffemodel")
            )
            print("WeChatQRCode检测器初始化成功")
            return True
        except Exception as e:
            print(f"检测器初始化失败: {e}")
            print("请确保模型文件已正确下载并放置于models目录")
            print("将使用模拟检测器代替，以便程序可以继续运行")
            
            # 创建一个模拟检测器类
            class MockQRDetector:
                def __init__(self):
                    self.current_task = 1  # 当前任务编号
                    self.current_chunk = 1  # 当前块编号
                    self.total_chunks_per_task = 5  # 每个任务的总块数
                    self.detection_count = 0  # 检测次数，用于控制任务切换
                    self.task_completed = False  # 当前任务是否已完成
                
                def detectAndDecode(self, frame):
                    import json
                    import time
                    
                    # 每30次检测模拟一次二维码扫描
                    self.detection_count += 1
                    if self.detection_count % 30 != 0:
                        return [], []
                    
                    # 模拟数据
                    results = []
                    points = []
                    
                    # 如果当前任务已完成，模拟两种情况：
                    # 1. 有时候返回相同的task_hash，测试"任务已完成"提示
                    # 2. 有时候返回新的task_hash，测试新任务开始
                    if self.task_completed:
                        # 70%的概率返回已完成的任务，30%的概率开始新任务
                        if self.detection_count % 10 < 7:
                            # 返回已完成的任务
                            task_hash = f"task_{self.current_task}_hash"
                            data = {
                                "task_hash": task_hash,
                                "total_chunks": self.total_chunks_per_task,
                                "current_chunk": 1  # 返回任意块，都会提示任务已完成
                            }
                            results.append(json.dumps(data))
                            points.append([[0,0], [0,100], [100,100], [100,0]])
                            print(f"模拟扫描到已完成的任务: {task_hash}")
                        else:
                            # 开始新任务
                            self.current_task += 1
                            self.current_chunk = 1
                            self.task_completed = False
                            task_hash = f"task_{self.current_task}_hash"
                            data = {
                                "task_hash": task_hash,
                                "total_chunks": self.total_chunks_per_task,
                                "current_chunk": self.current_chunk
                            }
                            results.append(json.dumps(data))
                            points.append([[0,0], [0,100], [100,100], [100,0]])
                            print(f"模拟扫描到新任务: {task_hash}")
                            self.current_chunk += 1
                    else:
                        # 返回当前任务的当前块
                        task_hash = f"task_{self.current_task}_hash"
                        data = {
                            "task_hash": task_hash,
                            "total_chunks": self.total_chunks_per_task,
                            "current_chunk": self.current_chunk,
                            "data": f"模拟数据块 {self.current_chunk}"
                        }
                        results.append(json.dumps(data))
                        points.append([[0,0], [0,100], [100,100], [100,0]])
                        print(f"模拟扫描到任务块: {task_hash} - {self.current_chunk}/{self.total_chunks_per_task}")
                        
                        # 更新块索引，判断任务是否完成
                        self.current_chunk += 1
                        if self.current_chunk > self.total_chunks_per_task:
                            self.task_completed = True
                            print(f"任务 {task_hash} 所有块已扫描完成")
                    
                    return results, points
            
            # 使用模拟检测器
            self.detector = MockQRDetector()
            return False  # 虽然返回False表示初始化失败，但提供了模拟检测器

    def set_work_time(self, start_hour: int, start_minute: int, end_hour: int, end_minute: int) -> None:
        """
        设置摄像头扫码的工作时间段
        
        参数:
            start_hour: 工作开始小时（0-23）
            start_minute: 工作开始分钟（0-59）
            end_hour: 工作结束小时（0-23）
            end_minute: 工作结束分钟（0-59）
        """
        # 验证输入参数
        if not (0 <= start_hour <= 23 and 0 <= start_minute <= 59 and 0 <= end_hour <= 23 and 0 <= end_minute <= 59):
            raise ValueError("时间参数必须在有效范围内（小时0-23，分钟0-59）")
        
        # 设置工作时间段
        self.work_start_time = (start_hour, start_minute)
        self.work_end_time = (end_hour, end_minute)
        print(f"工作时间段已设置为: {start_hour:02d}:{start_minute:02d} - {end_hour:02d}:{end_minute:02d}")
    
    def is_currently_in_work_time(self) -> bool:
        """
        检查当前时间是否在设定的工作时间段内
        
        返回:
            bool: 如果当前时间在工作时间段内返回True，否则返回False
        """
        # 获取当前时间
        current_time = time.localtime()
        current_hour = current_time.tm_hour
        current_minute = current_time.tm_min
        
        # 转换为分钟数进行比较
        current_total_minutes = current_hour * 60 + current_minute
        start_total_minutes = self.work_start_time[0] * 60 + self.work_start_time[1]
        end_total_minutes = self.work_end_time[0] * 60 + self.work_end_time[1]
        
        # 判断是否在工作时间内（考虑跨天情况）
        if start_total_minutes <= end_total_minutes:
            # 同一天内（如8:00-18:00）
            return start_total_minutes <= current_total_minutes <= end_total_minutes
        else:
            # 跨天情况（如20:00-06:00）
            return current_total_minutes >= start_total_minutes or current_total_minutes <= end_total_minutes
    
    def enable_work_time_check(self, enable: bool = True) -> None:
        """
        启用或禁用工作时间检查
        
        参数:
            enable: True表示启用工作时间检查（只在工作时间内工作），
                   False表示禁用工作时间检查（全天工作）
        """
        self.is_in_work_time = enable
        status = "已启用" if enable else "已禁用"
        print(f"工作时间检查功能{status}")
    
    def scan_qr_code_from_camera(self, capture_interval: float = 1.0) -> Optional[Dict]:
        """
        从摄像头按固定时间间隔拍摄并扫描二维码
        支持工作时间段检查功能
        
        参数:
            capture_interval: 拍摄间隔时间（秒），默认值为1秒
            
        返回:
            解码的数据字典，如果无法打开摄像头则返回None
        """        
        # 打开默认摄像头(0)
        cap = cv2.VideoCapture(0)
        
        # 检查摄像头是否成功打开
        if not cap.isOpened():
            print("无法打开摄像头")
            return None
        
        # 初始化变量
        decoded_data = []  # 存储所有解码的数据
        total_chunks = 100  # 总数据块数
        decoded_chunks = []  # 已解码的块索引列表
        task_hash = ""  # 任务哈希值，用于标识同一任务的二维码
        task_completed = False  # 任务完成标志
        first_scan = True  # 首次扫描标志
        last_capture_time = time.time()  # 上次拍摄时间
        
        print(f"开始按固定间隔拍摄，间隔时间: {capture_interval}秒")
        
        # 摄像头定时扫描循环
        while True:
            # 检查是否启用了工作时间检查
            if self.is_in_work_time:
                # 检查当前时间是否在工作时间段内
                if not self.is_currently_in_work_time():
                    current_time_str = time.strftime('%H:%M:%S')
                    start_time_str = f"{self.work_start_time[0]:02d}:{self.work_start_time[1]:02d}"
                    end_time_str = f"{self.work_end_time[0]:02d}:{self.work_end_time[1]:02d}"
                    
                    # 创建工作时间外的提示窗口
                    if 'frame' not in locals():
                        frame = np.zeros((480, 640, 3), dtype=np.uint8)
                    cv2.putText(frame, f'当前非工作时间: {current_time_str}', 
                                (50, 200), cv2.FONT_HERSHEY_SIMPLEX, 
                                1, (0, 0, 255), 2, cv2.LINE_AA)
                    cv2.putText(frame, f'工作时间: {start_time_str} - {end_time_str}', 
                                (50, 250), cv2.FONT_HERSHEY_SIMPLEX, 
                                1, (0, 0, 255), 2, cv2.LINE_AA)
                    cv2.putText(frame, '系统将在工作时间自动恢复...', 
                                (50, 300), cv2.FONT_HERSHEY_SIMPLEX, 
                                1, (0, 0, 255), 2, cv2.LINE_AA)
                    cv2.imshow('QR Code Scanner - Press Q to Quit', frame)
                    
                    # 检查键盘输入
                    key = cv2.waitKey(60000) & 0xFF  # 5秒检查一次，减少CPU占用
                    if key == ord('q'):  # 按Q键退出
                        break
                    continue  # 非工作时间跳过后续处理
            
            # 计算当前时间
            current_time = time.time()
            
            # 检查是否到达拍摄时间间隔
            if current_time - last_capture_time >= capture_interval:
                # 记录本次拍摄时间
                last_capture_time = current_time
                print(f"[{time.strftime('%H:%M:%S')}] 拍摄画面并扫描二维码...")
                
                # 在时间间隔到达时才执行cap.read()
                ret, frame = cap.read()
                if not ret:
                    print("无法读取摄像头画面")
                    break
                
                # 显示当前帧
                cv2.imshow('QR Code Scanner - Press Q to Quit', frame)
                
                # 扫描当前帧中的二维码
                results, points = self.detector.detectAndDecode(frame)

            # 只有在成功读取到图像和扫描结果时才处理二维码数据
            if 'results' in locals() and len(results) > 0:
                print("扫描到二维码：", results)

                # 处理扫描结果
                for i, result in enumerate(results):
                    try:
                        # 解析JSON数据
                        data_dict = json.loads(result)
                        current_task_hash = data_dict['task_hash']
                        current_total_chunks = data_dict['total_chunks']
                        
                        # 首次扫描或遇到新任务
                        if first_scan:
                            first_scan = False
                            task_hash = current_task_hash
                            total_chunks = current_total_chunks
                            decoded_chunks = []
                            decoded_data = []
                            task_completed = False
                            print(f"开始新任务：{task_hash}")
                        # 检查是否是同一任务的二维码
                        elif task_hash == current_task_hash:
                            # 如果任务已完成，提示用户
                            if task_completed:
                                print(f"此任务已完成，等待下一个扫码任务")
                                continue
                                
                            # 如果还有未解码的块
                            if len(decoded_chunks) < total_chunks:
                                # 如果当前块未被解码过
                                if data_dict['current_chunk'] not in decoded_chunks:
                                    # 记录已解码的块索引
                                    decoded_chunks.append(data_dict['current_chunk'])
                                    print("task:", task_hash, "_", data_dict['current_chunk'], " / ", total_chunks)
                                    # 添加二维码索引信息
                                    data_dict['qr_index'] = i
                                    # 保存解码数据
                                    decoded_data.append(data_dict)
                                    
                                    # 检查是否所有块都已解码
                                    if len(decoded_chunks) == total_chunks:
                                        # 重建数据
                                        df = self.reconstruct_dataframe(decoded_data)
                                        # 计算哈希值并验证
                                        hash_value = hash_dataframe(df)
                                        if hash_value == task_hash:
                                            print("重建准确，任务完成！！！！！！！！")
                                            task_completed = True
                                            
                                            # 将重建的数据保存到MongoDB
                                            # 使用任务哈希值作为集合名称的一部分，确保唯一性
                                            collection_name = f"qr_data_{task_hash[:10]}"
                                            print(f"尝试将数据保存到数据库，集合名: {collection_name}")
                                            save_success = self.insert_dataframe_to_mongo(df, collection_name)
                                            if save_success:
                                                print("数据已成功保存到MongoDB")
                                            else:
                                                print("数据保存到MongoDB失败")
                            else:
                                # 如果任务未标记为完成但所有块都已处理，标记为完成
                                if not task_completed:
                                    task_completed = True
                                    print("当前任务已完成！")
                        else:  # 发现新任务
                            print(f"发现新任务码：{current_task_hash}")
                            # 重置状态，准备新任务
                            task_hash = current_task_hash
                            total_chunks = current_total_chunks
                            decoded_chunks = []
                            decoded_data = []
                            task_completed = False
                            print("开始新的扫码任务")
                            
                            # 立即处理新任务的当前块
                            decoded_chunks.append(data_dict['current_chunk'])
                            print("task:", task_hash, "_", data_dict['current_chunk'], " / ", total_chunks)
                            data_dict['qr_index'] = i
                            decoded_data.append(data_dict)
                    except json.JSONDecodeError:
                        print(f"二维码数据不是有效的JSON格式")
                # 每次处理完结果后清理results变量
                del results
            
            # 检查键盘输入（使用较短的等待时间以保证响应速度）
            key = cv2.waitKey(100) & 0xFF
            if key == ord('q'):  # 按Q键退出
                break
            
            # 在非拍摄间隔期间短暂休眠，减少CPU占用
            if 'results' not in locals() or ('results' in locals() and len(results) == 0):
                time.sleep(0.05)
        
        # 释放摄像头资源
        cap.release()
        # 关闭所有OpenCV窗口
        cv2.destroyAllWindows()
        return None
    
    def reconstruct_dataframe(self, decoded_data: List[Dict]) -> pd.DataFrame:
        """
        根据解码的数据重建DataFrame
        
        参数:
            decoded_data: 解码的数据列表
            
        返回:
            重建后的完整DataFrame
            
        异常:
            ValueError: 如果没有可用的解码数据
        """
        # 检查是否有解码数据
        if not decoded_data:
            raise ValueError("没有可用的解码数据")
        
        # 按任务分块ID组织数据
        tasks: Dict[str, List[Dict]] = {}
        for data in decoded_data:
            current_chunk = str(data.get('current_chunk'))
            if current_chunk not in tasks:
                tasks[current_chunk] = []
            tasks[current_chunk].append(data)
        
        # 获取任务总块数
        total_chunks = decoded_data[0].get('total_chunks')
        
        # 检查数据完整性
        if len(decoded_data) != total_chunks:
            print(f"警告: 任务数据不完整，找到 {len(decoded_data)}/{total_chunks} 个分块")
        
        # 获取任务信息
        task_hash = decoded_data[0].get('task_hash')  # 修正拼写错误：tash_hash -> task_hash
        total_rows = decoded_data[0].get('total_rows')
        
        # 存储所有数据块
        all_chunks = []
        
        # 按顺序处理每个数据块
        for current_chunk in range(1, total_chunks + 1):
            # 获取当前块的数据
            chunk_data = tasks.get(str(current_chunk))[0]
            
            # 解析CSV数据
            csv_data = chunk_data.get('data', '')
            
            # 从字符串IO创建DataFrame
            chunk_df = pd.read_csv(StringIO(csv_data))
            
            # 添加到块列表
            all_chunks.append(chunk_df)

        # 如果有数据块，则合并
        if all_chunks:
            # 合并所有分块，忽略原索引
            full_df = pd.concat(all_chunks, ignore_index=True)
            
            # 验证重建后的数据行数
            if len(full_df) != total_rows:
                print(f"警告: 重建后行数 ({len(full_df)}) 与预期行数 ({total_rows}) 不一致")
            
            print(f"任务 {task_hash} 重建完成，实际行数: {len(full_df)}")
            return full_df
        else:
            # 如果没有数据块，返回空DataFrame
            return pd.DataFrame()

    def insert_dataframe_to_mongo(self, df: pd.DataFrame, collection_name: str, db_name: Optional[str] = None, del_existing: bool = False) -> bool:
        """
        将DataFrame数据插入到MongoDB数据库中
        
        参数:
            df: 要插入的pandas DataFrame
            collection_name: MongoDB集合名称
            db_name: 数据库名称，默认为None使用self.SVR_DB_NAME
            del_existing: 是否先删除集合中的所有现有数据，默认为False
            
        返回:
            bool: 操作成功返回True，失败返回False
        """
        if db_name is None:
            db_name = self.SVR_DB_NAME
            
        client = None
        try:
            # 检查DataFrame是否为空
            if df is None or df.empty:
                print("警告: 提供的DataFrame为空，无法插入数据库")
                return False
            
            # 创建MongoDB客户端连接
            client = MongoClient(self.SVR_MONGO_URI, serverSelectionTimeoutMS=5000)  # 添加超时设置
            # 验证连接
            client.server_info()
            
            db = client[db_name]
            coll = db[collection_name]
            
            # 如果需要，先删除集合中的所有数据
            if del_existing:
                coll.delete_many({})
                print(f"已删除集合 {collection_name} 中的所有现有数据")
            
            # 将DataFrame转换为字典列表并插入
            records = df.to_dict('records')
            coll.insert_many(records)
            
            # 添加插入时间戳字段（如果DataFrame中没有）
            if 'insert_time' not in df.columns:
                coll.update_many({}, {'$set': {'insert_time': time.time()}})
            
            print(f"成功将 {len(records)} 条记录插入到 {db_name}.{collection_name}")
            return True
            
        except Exception as e:
            print(f"将DataFrame插入MongoDB时出错: {str(e)}")
            return False
        finally:
            # 确保连接关闭
            if client is not None:
                try:
                    client.close()
                except:
                    pass
    
    def validate_reconstruction(self, original_df: pd.DataFrame, reconstructed_df: pd.DataFrame) -> bool:
        """
        验证数据重建的完整性
        
        参数:
            original_df: 原始DataFrame
            reconstructed_df: 重建的DataFrame
            
        返回:
            验证是否通过，完全匹配返回True，否则返回False
        """
        try:
            # 1. 检查行数是否匹配
            if len(original_df) != len(reconstructed_df):
                print(f"行数不匹配: 原始 {len(original_df)} vs 重建 {len(reconstructed_df)}")
                return False
            
            # 2. 检查列数是否匹配
            if len(original_df.columns) != len(reconstructed_df.columns):
                print(f"列数不匹配: 原始 {len(original_df.columns)} vs 重建 {len(reconstructed_df.columns)}")
                return False
            
            # 3. 检查列名是否完全一致
            if list(original_df.columns) != list(reconstructed_df.columns):
                print("列名不匹配")
                print(f"原始列: {list(original_df.columns)}")
                print(f"重建列: {list(reconstructed_df.columns)}")
                return False
            
            # 4. 检查数据内容（忽略NaN值的比较）
            # 重置索引以便正确比较
            original_reset = original_df.reset_index(drop=True)
            reconstructed_reset = reconstructed_df.reset_index(drop=True)
            
            # 先检查数据类型是否一致
            for col in original_reset.columns:
                if original_reset[col].dtype != reconstructed_reset[col].dtype:
                    print(f"列'{col}'的数据类型不匹配: 原始 {original_reset[col].dtype} vs 重建 {reconstructed_reset[col].dtype}")
                    # 尝试类型转换
                    try:
                        reconstructed_reset[col] = reconstructed_reset[col].astype(original_reset[col].dtype)
                        print(f"已将重建数据中'{col}'列转换为正确类型")
                    except:
                        print(f"无法转换列'{col}'的数据类型")
                        return False
            
            # 使用pandas的equals方法进行高效比较（已处理NaN值）
            # 这比逐单元格比较更高效且更可靠
            if original_reset.equals(reconstructed_reset):
                print("数据完整性验证通过!")
                return True
            else:
                # 如果不匹配，尝试找出第一个不匹配的位置
                for col in original_reset.columns:
                    for i in range(len(original_reset)):
                        orig_val = original_reset.at[i, col]
                        recon_val = reconstructed_reset.at[i, col]
                        
                        # 特殊处理NaN值比较
                        if pd.isna(orig_val) and pd.isna(recon_val):
                            continue  # 两个都是NaN，认为相等
                        # 其他情况直接比较值
                        elif orig_val != recon_val:
                            print(f"数据不匹配: 第{i}行, 列'{col}': 原始={orig_val}({type(orig_val)}) vs 重建={recon_val}({type(recon_val)})")
                            return False
            return False  # 以防万一
            
        except Exception as e:
            # 捕获比较过程中可能出现的异常
            print(f"数据比较时出错: {e}")
            import traceback
            traceback.print_exc()
            return False

if __name__ == "__main__":
    try:
        # 创建二维码数据重建器实例
        reconstructor = QRCodeDataReconstructor()
        
        # 设置工作时间段（例如：早8点到晚6点）
        reconstructor.set_work_time(start_hour=8, start_minute=0, end_hour=18, end_minute=0)
        
        # 启用工作时间检查
        reconstructor.enable_work_time_check(enable=True)
        
        # 启动摄像头扫描二维码，设置拍摄间隔为1秒
        # 可以根据需要调整间隔时间，例如: scan_qr_code_from_camera(capture_interval=2.0)
        result = reconstructor.scan_qr_code_from_camera(capture_interval=1.0)
        print(f"程序结束，结果: {result}")
    except KeyboardInterrupt:
        print("\n程序被用户中断")
    except Exception as e:
        print(f"程序运行出错: {e}")
        import traceback
        traceback.print_exc()
    finally:
        # 确保释放所有资源
        cv2.destroyAllWindows()  # 确保所有OpenCV窗口都被关闭